import os.path

from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector
from mmcv import Config
import mmcv
import os.path as osp
from mmdet.apis import set_random_seed
import time
import argparse


def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='train config file path')
    args = parser.parse_args()
    return args


def main(cfg_path):
    # try:
    assert cfg_path is not None and os.path.exists(cfg_path), 'cfg_path不能为none'
    cfg = Config.fromfile(cfg_path)
    # Build dataset
    cfg.gpu_ids = range(1)
    cfg.seed = 0
    set_random_seed(0, deterministic=False)
    datasets = [build_dataset(cfg.data.train)]

    # Build the detector
    model = build_detector(cfg.model)
    # Add an attribute for visualization convenience
    model.CLASSES = datasets[0].CLASSES

    # Create work_dir
    mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
    cfg.dump(osp.join(cfg.work_dir, osp.basename(cfg_path)))
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    train_detector(model, datasets, cfg, distributed=False, validate=True, timestamp=timestamp)
    # except Exception:
    #     with open('./error_logs.txt', mode='a+', encoding='utf-8') as f:
    #         f.write(time.strftime('%Y-%m-%d_%H:%M:%S', time.localtime()) + '\n')
    #         f.write(cfg_path + '\n')
    #         # f.write(Exception.__name__ + '\n')

if __name__ == '__main__':
    main(parse_args().config)
